package com.liru.neo4j_demo.config;

import org.neo4j.driver.Driver;
import org.neo4j.driver.Session;
import org.neo4j.driver.Transaction;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;
import org.springframework.stereotype.Component;

/**
 * Neo4j向量索引初始化器
 * 应用启动时创建必要的向量索引
 */
@Component
public class Neo4jIndexInitializer implements CommandLineRunner {

    private final Driver neo4jDriver;

    @Autowired
    public Neo4jIndexInitializer(Driver neo4jDriver) {
        this.neo4jDriver = neo4jDriver;
    }

    @Override
    public void run(String... args) throws Exception {
        createVectorIndex();
    }

    /**
     * 创建向量索引
     */
    private void createVectorIndex() {
        try (Session session = neo4jDriver.session()) {
            try (Transaction tx = session.beginTransaction()) {
                // 创建Chunk节点的embedding向量索引
                tx.run(
                        "CREATE VECTOR INDEX chunk_embedding_idx " +
                        "FOR (c:Chunk) ON (c.embedding) " +
                        "OPTIONS { " +
                        "  indexConfig: { " +
                        "    `vector.dimensions`: 384, " +
                        "    `vector.similarity_function`: 'cosine' " +
                        "  } " +
                        "}"
                );

                tx.commit();
                System.out.println("Neo4j向量索引创建成功");
            } catch (Exception e) {
                System.out.println("Neo4j向量索引创建失败: " + e.getMessage());
                // 如果索引已存在，忽略错误
            }
        }
    }
}